package com.solon.open.service;

import com.google.common.collect.Lists;
import com.hankcs.hanlp.seg.common.Term;
import com.hankcs.hanlp.tokenizer.NLPTokenizer;
import com.solon.open.config.QaConst;
import com.solon.open.entity.KeyWord;
import com.solon.open.mapper.QaMilvusMapper;
import com.solon.open.model.QaModel;
import org.apache.commons.collections4.CollectionUtils;
import org.dromara.milvus.plus.model.vo.MilvusResp;
import org.dromara.milvus.plus.model.vo.MilvusResult;
import org.noear.solon.annotation.Component;
import org.noear.solon.annotation.Inject;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

@Component
public class QaService {
    @Inject
    QaMilvusMapper qaMilvusMapper;

    public void addQa(List<QaModel> model){
        for (QaModel qaModel : model) {
            List<String> keyword = keyword(qaModel.getQuestion());
            qaModel.setKeyword(new KeyWord(keyword));
        }
        qaMilvusMapper.insert(model);
    }
    public String getQa(String question,List<Float> v){
        List<String> keyword = keyword(question);
        MilvusResp<List<MilvusResult<QaModel>>> query = qaMilvusMapper.
                 queryWrapper().vector(v)
                .jsonContainsAny("keyword[\"kw\"]",keyword)
                .topK(1).query();
        List<MilvusResult<QaModel>> data = query.getData();
        if(CollectionUtils.isNotEmpty(data)){
            MilvusResult<QaModel> qaModelMilvusResult = data.get(0);
            if(qaModelMilvusResult.getDistance()<0.3){
                return qaModelMilvusResult.getEntity().getAnswer();
            }
        }
        return QaConst.DEFAULT_ANSWER;
    }
    private List<String> keyword(String inputStr) {
        ArrayList<String> natures = Lists.newArrayList("n","Ng","v");
        List<Term> segment = NLPTokenizer.segment(inputStr);
        List<String> keyword = segment.stream().filter(v -> natures.contains(v.nature.toString())).map(a -> a.word).collect(Collectors.toList());
        return keyword;
    }
}
